﻿using System;
/*	http://code.google.com/p/wamon/
 *	Windows Activity Monitor - Monitors user's activity on the computer, 
 *	tracks active windows and processes, allows table and chart visualization.
 *	Copyright (C) 2009  Archae s.r.o. (http://archaedev.com/)
 *
 *	This program is free software: you can redistribute it and/or modify
 *	it under the terms of the GNU General Public License as published by
 *	the Free Software Foundation, either version 3 of the License, or
 *	(at your option) any later version.
 *
 *	This program is distributed in the hope that it will be useful,
 *	but WITHOUT ANY WARRANTY; without even the implied warranty of
 *	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *	GNU General Public License for more details.
 *
 *	You should have received a copy of the GNU General Public License
 *	along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
using System.Threading;
using System.Runtime.InteropServices;
using System.Text;
using System.Collections.Generic;

namespace Archae.Wamon
{
	class WamonDaemon
	{
		private Thread thread;

		public void Start()
		{
			if (thread != null)
				Stop();
			thread = new Thread(new ThreadStart(Loop));
			thread.Start();
		}

		public void Stop()
		{
			if (thread != null)
			{
				thread.Abort();
				thread.Join();
				thread = null;
			}
		}

		private void Loop()
		{
			LASTINPUTINFO lastInputInfo = new LASTINPUTINFO();
			lastInputInfo.cbSize = (uint)Marshal.SizeOf(lastInputInfo);
			lastInputInfo.dwTime = 0;
			uint[] processes = new uint[1000];
			try
			{
				while (true)
				{
					Thread.Sleep(60000);
					try
					{
						bool isIdle = (GetLastInputInfo(ref lastInputInfo) && Environment.TickCount - lastInputInfo.dwTime > 60000);
						if (isIdle && DB.Singleton.RecordIdle)
						{
							string userName = GetUserName();
							if (userName == "SYSTEM")
								throw new Exception("Ignoring system window");
							DB.Singleton.InsertSnapshot("IDLE", "IDLE", userName);
						}
						else
						{
							IntPtr hwnd = GetForegroundWindow();
							if (hwnd == IntPtr.Zero || !IsWindow(hwnd))
								throw new Exception("No window in foreground");
							string windowText = GetWindowText(hwnd);
							string process, user;
							uint processId;
							GetWindowThreadProcessId(hwnd, out processId);
							GetProcessInfo(processId, out process, out user);
							if (string.IsNullOrEmpty(windowText) || string.IsNullOrEmpty(process) || user == "SYSTEM")
								throw new Exception("Ignoring empty window, process or system user");
							DB.Singleton.InsertSnapshot(windowText, process, user);
						}
					}
					catch (ThreadAbortException) { }
					catch (Exception e) { DB.Singleton.Log(e.Message); }
					try
					{
						List<string> insertedProcesses = new List<string>();
						uint copied = 0;
						if (EnumProcesses(processes, 1000, out copied))
						{
							for (int i = 0; i < copied; i++)
							{
								string process, user;
								GetProcessInfo(processes[i], out process, out user);
								if (string.IsNullOrEmpty(process) || process.IndexOf('\\') == -1 || insertedProcesses.Contains(process))
									continue;
								uint id = DB.Singleton.GetProcessID(process.Substring(process.LastIndexOf('\\') + 1));
								if (id == 0)
									continue;
								DB.Singleton.InsertProcessSnapshot(id, user);
							}
						}
					}
					catch (ThreadAbortException) { }
					catch (Exception e) { DB.Singleton.Log(e.Message); }
				}
			}
			catch (ThreadAbortException) { }
		}

		private string GetUserName()
		{
			StringBuilder buffer = new StringBuilder(256);
			int size = 256;
			GetUserName(buffer, ref size);
			return buffer.ToString();
		}

		private string GetWindowText(IntPtr hwnd)
		{
			int length = GetWindowTextLength(hwnd);
			StringBuilder buffer = new StringBuilder(length + 1);
			GetWindowText(hwnd, buffer, buffer.Capacity);
			return buffer.ToString();
		}

		private void GetProcessInfo(uint processId, out string process, out string user)
		{
			process = user = null;
			// Open the process handle
			IntPtr hProcess = OpenProcess(ProcessAccessFlags.QueryInformation | ProcessAccessFlags.VMRead, false, processId);
			if (hProcess == IntPtr.Zero)
				return;
			// Get module file name
			StringBuilder buffer = new StringBuilder(4096);
			GetModuleFileNameEx(hProcess, IntPtr.Zero, buffer, buffer.Capacity);
			process = buffer.ToString();
			// Get process user name
			IntPtr hToken;
			if (OpenProcessToken(hProcess, 0x0008, out hToken))
			{
				uint len = 0;
				GetTokenInformation(hToken, TOKEN_INFORMATION_CLASS.TokenUser, IntPtr.Zero, len, out len);
				IntPtr tokenInformation = Marshal.AllocHGlobal((int)len);
				if (GetTokenInformation(hToken, TOKEN_INFORMATION_CLASS.TokenUser, tokenInformation, len, out len))
				{
					TOKEN_USER tokenUser = (TOKEN_USER)Marshal.PtrToStructure(tokenInformation, typeof(TOKEN_USER));
					uint sidLength = GetLengthSid(tokenUser.User.Sid);
					byte[] sid = new byte[sidLength];
					Marshal.Copy(tokenUser.User.Sid, sid, 0, (int)sidLength);
					len = 4096;
					buffer = new StringBuilder((int)len);
					uint len2 = 300;
					StringBuilder buffer2 = new StringBuilder((int)len2);
					SID_NAME_USE use;
					if (LookupAccountSid(null, sid, buffer, ref len, buffer2, ref len2, out use))
						user = buffer.ToString();
				}
				CloseHandle(hToken);
			}
			// Close the process handle
			CloseHandle(hProcess);
			// Make sure there is some user name
			if (string.IsNullOrEmpty(user))
				user = GetUserName();
		}

		#region P/Invoke Methods
		[StructLayout(LayoutKind.Sequential)]
		private struct LASTINPUTINFO
		{
			public static readonly int SizeOf = Marshal.SizeOf(typeof(LASTINPUTINFO));
			[MarshalAs(UnmanagedType.U4)]
			public UInt32 cbSize;
			[MarshalAs(UnmanagedType.U4)]
			public UInt32 dwTime;
		}
		[DllImport("user32.dll")]
		private static extern bool GetLastInputInfo(ref LASTINPUTINFO plii);
		[DllImport("Advapi32.dll")]
		private static extern bool GetUserName(StringBuilder lpBuffer, ref int nSize);
		[DllImport("user32.dll")]
		private static extern IntPtr GetForegroundWindow();
		[DllImport("user32.dll")]
		[return: MarshalAs(UnmanagedType.Bool)]
		private static extern bool IsWindow(IntPtr hWnd);
		[DllImport("user32.dll", CharSet = CharSet.Auto, SetLastError = true)]
		private static extern int GetWindowText(IntPtr hWnd, StringBuilder lpString, int nMaxCount);
		[DllImport("user32.dll", SetLastError = true, CharSet = CharSet.Auto)]
		private static extern int GetWindowTextLength(IntPtr hWnd);
		[DllImport("user32.dll", SetLastError = true)]
		private static extern uint GetWindowThreadProcessId(IntPtr hWnd, out uint lpdwProcessId);
		[Flags]
		private enum ProcessAccessFlags : uint
		{
			All = 0x001F0FFF,
			Terminate = 0x00000001,
			CreateThread = 0x00000002,
			VMOperation = 0x00000008,
			VMRead = 0x00000010,
			VMWrite = 0x00000020,
			DupHandle = 0x00000040,
			SetInformation = 0x00000200,
			QueryInformation = 0x00000400,
			Synchronize = 0x00100000
		}
		[DllImport("kernel32.dll")]
		private static extern IntPtr OpenProcess(ProcessAccessFlags dwDesiredAccess, [MarshalAs(UnmanagedType.Bool)] bool bInheritHandle, uint dwProcessId);
		[DllImport("psapi.dll", SetLastError = true)]
		private static extern int GetModuleFileNameEx(IntPtr hProcess, IntPtr hModule, StringBuilder lpFilename, int nSize);
		[DllImport("kernel32.dll", SetLastError = true)]
		[return: MarshalAs(UnmanagedType.Bool)]
		private static extern bool CloseHandle(IntPtr hObject);
		[DllImport("advapi32.dll", SetLastError = true)]
		[return: MarshalAs(UnmanagedType.Bool)]
		private static extern bool OpenProcessToken(IntPtr ProcessHandle, UInt32 DesiredAccess, out IntPtr TokenHandle);
		private enum TOKEN_INFORMATION_CLASS
		{
			TokenUser = 1,
			TokenGroups,
			TokenPrivileges,
			TokenOwner,
			TokenPrimaryGroup,
			TokenDefaultDacl,
			TokenSource,
			TokenType,
			TokenImpersonationLevel,
			TokenStatistics,
			TokenRestrictedSids,
			TokenSessionId,
			TokenGroupsAndPrivileges,
			TokenSessionReference,
			TokenSandBoxInert,
			TokenAuditPolicy,
			TokenOrigin,
			TokenElevationType,
			TokenLinkedToken,
			TokenElevation,
			TokenHasRestrictions,
			TokenAccessInformation,
			TokenVirtualizationAllowed,
			TokenVirtualizationEnabled,
			TokenIntegrityLevel,
			TokenUIAccess,
			TokenMandatoryPolicy,
			TokenLogonSid,
			MaxTokenInfoClass
		}
		public struct TOKEN_USER
		{
			public SID_AND_ATTRIBUTES User;
		}
		[StructLayout(LayoutKind.Sequential)]
		public struct SID_AND_ATTRIBUTES
		{
			public IntPtr Sid;
			public int Attributes;
		}
		[DllImport("advapi32.dll", SetLastError = true)]
		private static extern bool GetTokenInformation(IntPtr TokenHandle, TOKEN_INFORMATION_CLASS TokenInformationClass, IntPtr TokenInformation, uint TokenInformationLength, out uint ReturnLength);
		enum SID_NAME_USE
		{
			SidTypeUser = 1,
			SidTypeGroup,
			SidTypeDomain,
			SidTypeAlias,
			SidTypeWellKnownGroup,
			SidTypeDeletedAccount,
			SidTypeInvalid,
			SidTypeUnknown,
			SidTypeComputer
		}
		[DllImport("advapi32.dll", CharSet = CharSet.Auto, SetLastError = true)]
		private static extern bool LookupAccountSid(string lpSystemName, [MarshalAs(UnmanagedType.LPArray)] byte[] Sid, StringBuilder lpName, ref uint cchName, 
			StringBuilder ReferencedDomainName, ref uint cchReferencedDomainName, out SID_NAME_USE peUse);
		[DllImport("advapi32.dll")]
		private static extern uint GetLengthSid(IntPtr sid);
		[DllImport("psapi.dll", SetLastError = true)]
		static extern bool EnumProcesses([MarshalAs(UnmanagedType.LPArray)] uint[] processIds, uint arraySizeBytes, out uint bytesCopied);
		#endregion
	}
}
